// src/lib.rs
use libc::{c_float, c_uint};
use rustfft::{num_complex::{Complex32, ComplexFloat}, FftDirection, FftPlanner};
use lazy_static::lazy_static;
use std::sync::Mutex;

// 全局 FftPlanner 单例（用 Mutex 保证线程安全）
lazy_static! {
    static ref FFT_PLANNER: Mutex<FftPlanner<f32>> = Mutex::new(FftPlanner::new());
}

#[deprecated(note = "use Complex type provided by rustfft")]
#[repr(C)]
pub struct Complex {
    re: c_float,
    im: c_float,
}

// 将Kotlin传递的指针转换为Rust的Complex32数组
fn from_raw(ptr: *const Complex, len: usize) -> Vec<Complex32> {
    let slice = unsafe { std::slice::from_raw_parts(ptr, len) };
    slice.iter().map(|c| Complex32::new(c.re, c.im)).collect()
}

// 将Rust的Complex32数组转换为C兼容结构
fn to_raw(buffer: &[Complex32]) -> Vec<Complex> {
    buffer
        .iter()
        .map(|c| Complex { re: c.re, im: c.im })
        .collect()
}

#[unsafe(no_mangle)]
pub extern "C" fn add(a: i32, b: i32) -> i32 {
    a + b
}

#[unsafe(no_mangle)]
#[deprecated(note = "use `try_fft` instead")]
pub extern "C" fn fft_forward(
    input: *const Complex,
    output: *mut Complex,
    n: c_uint,
    inverse: bool,
) {
    let len = n as usize;
    let mut input_buffer = from_raw(input, len);
    let mut planner = FftPlanner::new();
    let direction = if inverse {
        FftDirection::Inverse
    } else {
        FftDirection::Forward
    };
    let fft = planner.plan_fft(len, direction);
    fft.process(&mut input_buffer);

    // RustFFT does not normalize outputs. 
    if direction == FftDirection::Inverse {
        for c in &mut input_buffer {
            let n = len as f32;
            c.re = c.re / n;
            c.im = c.im / n;
        }
    }

    let output_buffer = to_raw(&input_buffer);
    unsafe {
        std::ptr::copy_nonoverlapping(output_buffer.as_ptr(), output, len);
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn try_fft(
    ptr: *mut Complex32,
    n: c_uint,
    inverse: bool,
) {
    let len = n as usize;
    let mut buffer = unsafe { std::slice::from_raw_parts_mut(ptr, len) };
    
    let direction = if inverse {
        FftDirection::Inverse
    } else {
        FftDirection::Forward
    };

    let mut planner = FFT_PLANNER.lock().unwrap(); // 多线程场景需处理锁竞争，或用 RwLock
    let fft = planner.plan_fft(len, direction);
    
    fft.process(&mut buffer);

    // RustFFT does not normalize outputs. 
    if direction == FftDirection::Inverse {
        let n = len as f32;
        for c in buffer {
            c.re = c.re / n;
            c.im = c.im / n;
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn magnitude32(input: *const Complex32, n: c_uint , output: *mut c_float) {
    let len = n as usize;
    let slice = unsafe {
        std::slice::from_raw_parts(input, len)
    };
    let output_slice = unsafe {
        std::slice::from_raw_parts_mut(output, len)
    };

    for i in 0..len {
        output_slice[i] = slice[i].abs();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_works() {
        let result = add(2, 2);
        assert_eq!(result, 4);
    }

    #[test]
    fn test_magnitude32() {
        let cpx = Complex32 {
            re: -4.,
            im: 3.
        };
        assert_eq!(cpx.abs(), 5.);
    }
}
